(* A small lisp interpreter.                            * 
 * Shamelessly stolen from http://norvig.com/lispy.html *)
exception IllFormedExpr of string
exception UndefinedVariable of string
exception NotImplemented of string

module EnvMap = Map.Make (struct type t = string let compare = compare end)
type env = { mutable env : Types.Data.t EnvMap.t;
	     outerEnv    : env option
	   }

let makeEnv2 vars args outerEnv =
  match vars, args with
    | Types.Data.List varl, Types.Data.List argl ->
      let env = 
	List.fold_left2 (fun a var arg -> EnvMap.add var arg a)
	  EnvMap.empty (List.map Types.Data.symbol_to_string varl) argl 
      in { env; outerEnv }
    | Types.Data.Symbol var, Types.Data.Symbol arg ->
      let env = EnvMap.add var args EnvMap.empty 
      in { env; outerEnv } 
    | _ -> raise (IllFormedExpr (Types.Data.show vars))

let rec findVar var env = 
  try EnvMap.find var env.env
  with Not_found -> 
    match env.outerEnv with
      | None -> raise (UndefinedVariable var)
      | Some outer -> findVar var outer

let globalEnv = (* set up the global environment *)
  let operator operInt operFloat acc = Types.Data.Fun 
    (function (* generic function for arithmetic operators *)
      | (Types.Data.List ls) -> 
      begin 
	if List.exists (function (Types.Data.Float a) -> true | _ -> false) ls
	then (* circumventing the Ocaml type system *)
	  let args = List.map (fun a -> match a with 
	    | Types.Data.Float af -> af
	    | Types.Data.Int ai -> float_of_int ai 
	    | _ -> raise (Invalid_argument "arithmetic")) ls in
	  match args with [] -> Types.Data.Float (float_of_int acc) | h::t ->
	    Types.Data.Float (List.fold_left operFloat h t)
	else
	  let args = List.map (fun a -> match a with 
	    | Types.Data.Int ai -> ai 
	    | _ -> raise (Invalid_argument "arithmetic")) ls in 
	  match args with [] -> Types.Data.Float (float_of_int acc) | h::t ->
	    Types.Data.Int (List.fold_left operInt h t)
      end
      | _ -> assert false)
  and eqOperator operInt operFloat acc = Types.Data.Fun 
    (function (* generic function for comparison operators *)
      | (Types.Data.List ls) -> 
      begin 
	if List.exists (function (Types.Data.Float a) -> true | _ -> false) ls
	then (* circumventing the Ocaml type system *)
	  let args = List.map (fun a -> match a with 
	    | Types.Data.Float af -> af
	    | Types.Data.Int ai -> float_of_int ai 
	    | _ -> raise (Invalid_argument "compare")) ls in
	  match args with [] -> Types.Data.Bool acc | h::t ->
	    Types.Data.Bool (List.exists (fun x -> operFloat h x) t)
	else
	  let args = List.map (fun a -> match a with 
	    | Types.Data.Int ai -> ai 
	    | _ -> raise (Invalid_argument "compare")) ls in 	  
	  match args with [] -> Types.Data.Bool acc | h::t ->
	    Types.Data.Bool (List.exists (fun x -> (operInt h x)) t)
      end
      | _ -> assert false)
  and logOperator oper acc = Types.Data.Fun 
    (function (* generic function for comparison operators *)
      | (Types.Data.List ls) -> let args = List.map Types.Data.to_bool ls in
	    Types.Data.Bool (List.fold_left (fun a x -> (oper a x)) acc args)
      | _ -> assert false)
  in
  let globals = (* some basic built-in functions *)
    [("+",   operator ( + ) ( +.) 0);
     ("-",   operator ( - ) ( -.) 0);
     ("*",   operator ( * ) ( *.) 1);
     ("/",   operator ( / ) ( /.) 1);
     (">", eqOperator (>  ) (>  ) true);
     ("<", eqOperator (<  ) (<  ) true);
     (">=",eqOperator (>= ) (>= ) true);
     ("<=",eqOperator (<= ) (<= ) true);
     ("=", eqOperator ( = ) ( = ) true);
     ("max", operator  max   max  0);
     ("min", operator  min   min  0);
     ("not",  Types.Data.Fun (function 
                       | (Types.Data.List ((Types.Data.Bool b)::[])) -> 
			    Types.Data.Bool (not b)
                       | _ -> raise (Invalid_argument "not")));
     ("or", logOperator (||) false);
     ("and",  logOperator (&&) true);
     ("modulo", Types.Data.Fun (function 
                       | (Types.Data.List 
			    ((Types.Data.Int a)::(Types.Data.Int b)::[])) -> 
			    Types.Data.Int (a mod b)
                       | _ -> raise (Invalid_argument "len")));
     ("length", Types.Data.Fun (function 
                       | (Types.Data.List ((Types.Data.List ls)::[])) -> 
			    Types.Data.Int (List.length ls)
                       | _ -> raise (Invalid_argument "len")));
     ("cons", Types.Data.Fun Types.Data.cons);
     ("car", Types.Data.Fun (function 
                       | Types.Data.List ((Types.Data.List ls)::[]) -> List.hd ls
                       | _ -> raise (Invalid_argument "car")));
     ("cdr", Types.Data.Fun (function 
                       | Types.Data.List ((Types.Data.List ls)::[]) -> 
			   Types.Data.List (List.tl ls)
                       | _ -> raise (Invalid_argument "cdr")));
     ("append", Types.Data.Fun Types.Data.append);
     ("list", Types.Data.Fun (function a -> a));
     ("list?", Types.Data.Fun (function 
                       | Types.Data.List ((Types.Data.List _)::[]) -> 
	                   Types.Data.Bool true 
                       | _ -> Types.Data.Bool false));
     ("null?", Types.Data.Fun (function 
                       | (Types.Data.List ((Types.Data.List [])::[]))-> 
			   Types.Data.Bool true
                       | _ -> Types.Data.Bool false));
     ("symbol?", Types.Data.Fun (function
                       | (Types.Data.List ((Types.Data.Symbol _)::[]))->
			   Types.Data.Bool true
                       | _ -> Types.Data.Bool false));
     ("sin", Types.Data.Fun (function
                       | (Types.Data.List ((Types.Data.Float ang)::[]))->
			   Types.Data.Float (sin ang)
                       | (Types.Data.List ((Types.Data.Int ang)::[]))->
			   Types.Data.Float (sin (float_of_int ang))
                       | _ -> raise (Invalid_argument "sin")));
     ("cos", Types.Data.Fun (function
                       | (Types.Data.List ((Types.Data.Float ang)::[]))->
			   Types.Data.Float (cos ang)
                       | (Types.Data.List ((Types.Data.Int ang)::[]))->
			   Types.Data.Float (cos (float_of_int ang))
                       | _ -> raise (Invalid_argument "sin")))] in
  let lists = (List.split globals) in
  makeEnv2 (Types.Data.of_list (List.map Types.Data.symbol_of_string (fst lists))) 
    (Types.Data.of_list (snd lists)) None

let addGlobal var newVal =
  let gEnv = globalEnv.env in
  globalEnv.env <- EnvMap.add var newVal gEnv

let setGlobal var newVal =
  if EnvMap.mem var globalEnv.env then addGlobal var newVal
  else raise (UndefinedVariable var)

let rec eval expression env =
  match expression with
    | Types.Data.Symbol a -> findVar a env (* variable reference *)
    | Types.Data.List ((Types.Data.Symbol func)::args) ->
      (match func with
	| "quote" -> (* (quote a) *)
	  begin
	    match args with
	      | a::[] -> a
	      | _ -> raise (IllFormedExpr (Types.Data.show expression))
	  end
	| "if" -> (* (if test caseT caseF) *)
	  begin 
	    match args with
	      | test::caseT::caseF::[] -> eval (match eval test env with
		  | Types.Data.Bool false -> caseF
		  | _ -> caseT) env
	      | _ -> raise (IllFormedExpr (Types.Data.show expression))
	  end
	| "set!" -> (* (set! var expr) *)
	  begin 
	    match args with
	      | (Types.Data.Symbol var)::expr::[] -> 
		  (setGlobal var (eval expr env); Types.Data.Unit)
	      | _ -> raise (IllFormedExpr (Types.Data.show expression))
	  end
	| "lambda" -> (* (lambda argn body) *)
	  begin 
	    match args with
	      | argn::body::[] -> Types.Data.Fun 
		  (function argl -> eval body (makeEnv2 argn argl (Some env)))
	      | _ -> raise (IllFormedExpr (Types.Data.show expression))
	  end
	| "define" -> (* (define var expr) *)
	  begin 
	    match args with
	      | var::expr::[] -> if env == globalEnv then
		  (env.env <- EnvMap.add (Types.Data.symbol_to_string var) 
		     (eval expr env) env.env; 
		   Types.Data.Unit) else Types.Data.Unit
	      | _ -> raise (IllFormedExpr (Types.Data.show expression))
	  end
	| "begin" -> List.fold_left (fun a b -> eval b env) Types.Data.Unit args
	| _ -> let evaled = List.map (fun x -> eval x env) (* a function *)
		 (Types.Data.to_list expression) in
	       begin
		 match evaled with
		   | (Types.Data.Fun f)::argl -> f (Types.Data.of_list argl)
		   | _ -> assert false
	       end)
    | Types.Data.List (f::args) -> (* the function name is not a symbol *)
        raise (IllFormedExpr ((Types.Data.show f)^" is not a function"))
    | Types.Data.List [] -> raise (IllFormedExpr "()") (* empty list () *)
    | _ -> expression

let repl (in_channel, out_channel) = (* a simple read-eval-print loop *)
  let () = Printf.printf "> " in
  let () = flush stdout in
  let rec loop () = 
    let () = 
      try
	let lexbuf = Lexing.from_channel in_channel in
	let expr = Parser.main Lexer.token lexbuf in
	let result = eval expr globalEnv in
	let () = Printf.fprintf out_channel "%s\n> " (Types.Data.show result) in
	let () = flush out_channel in
	()
      with Lexer.Eof -> ()
	| Parsing.Parse_error ->
	  let () = Printf.fprintf out_channel "Parse error\n> " in
	  let () = flush out_channel in ()
	| IllFormedExpr s -> 
	  let () = Printf.fprintf out_channel "Ill formed expression: %s\n> " s in
	  let () = flush out_channel in ()
	| UndefinedVariable s ->
	  let () = Printf.fprintf out_channel "Undefined variable: %s\n> " s in
	  let () = flush out_channel in ()
	| NotImplemented s ->
	  let () = Printf.fprintf out_channel "Not implemented: %s\n> " s in
	  let () = flush out_channel in ()
	| _ -> 
	  let () = Printf.fprintf out_channel "Error\n> " in
	  let () = flush out_channel in ()
    in loop ()
  in loop ()

